package com.szl.group.config.aop;

import com.google.common.collect.ImmutableList;
import com.szl.group.common.anno.Limit;
import com.szl.group.common.enums.LimitType;
import com.szl.group.exception.BusinessException;
import com.szl.group.utils.AddressIpUtils;
import com.szl.group.utils.JoinPointUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;

import static com.szl.group.common.enums.ResultCodeEnum.FREQUENTLY;


/**
 * 基于lua脚本实现分布式限流
 */
@Component
@Slf4j
@Aspect
@Order(2)
public class LimitInterceptor {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;


    @Around("execution(public * *(..)) && @annotation(com.szl.group.common.anno.Limit)")
    public Object interceptor(ProceedingJoinPoint pjp) {

        Method method = JoinPointUtils.getMethod(pjp);
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();

        String name = limitAnnotation.name();
        String key;

        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();

        switch (limitType) {
            case IP:
                HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
                key = AddressIpUtils.getIpAddress(request);
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }

        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(), key));

        try {

            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
            log.info("Access try count is {} for name={} and key = {}", count, name, key);

            if (count != null && count.intValue() <= limitCount) {
                return pjp.proceed();
            } else {
                throw new BusinessException(FREQUENTLY, "You have been dragged into the blacklist");
            }

        } catch (Throwable e) {
            if (e instanceof BusinessException) {
                throw new BusinessException(FREQUENTLY, e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
    }

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    public String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

    private static final String UNKNOWN = "unknown";





}
